package com.learning.rocketmq.optimal.throughput;

import org.apache.rocketmq.client.consumer.DefaultMQPullConsumer;
import org.apache.rocketmq.client.consumer.MQPullConsumer;
import org.apache.rocketmq.client.consumer.MQPullConsumerScheduleService;
import org.apache.rocketmq.client.consumer.MessageQueueListener;
import org.apache.rocketmq.client.consumer.PullResult;
import org.apache.rocketmq.client.consumer.PullTaskCallback;
import org.apache.rocketmq.client.consumer.PullTaskContext;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.message.MessageQueue;
import org.apache.rocketmq.common.protocol.heartbeat.MessageModel;

import java.util.Set;

/**
 * ClassName: PullConsumerRebalance
 * Description: DefaultMQPushConsumer 的负载均衡
 * <p>
 * `DefaultMQPullConsumer` 有两个辅助方法可以帮助实现负载均衡
 * - `registerMessageQueueListener()` 方法， 在有新的`Consumer` 加人或退出时被触发
 * - `MQPullConsumerScheduleService` 类，
 * Date: 2019/1/15 10:31 【需求编号】
 *
 * @author Sam Sho
 * @version V1.0.0
 */
public class PullConsumerRebalance {


    /**
     * `registerMessageQueueListener()` 方法， 在有新的`Consumer` 加人或退出时被触发
     *
     * @throws MQClientException
     */
    public static void consumer() throws MQClientException {
        DefaultMQPullConsumer consumer = new DefaultMQPullConsumer("please_rename_unique_group_name_5");

        consumer.start();

        consumer.registerMessageQueueListener("TopicTest1", new MessageQueueListener() {
            @Override
            public void messageQueueChanged(String topic, Set<MessageQueue> mqAll, Set<MessageQueue> mqDivided) {
                // 负载处理
            }
        });
    }


    /**
     * MQPullConsumerScheduleService 类的实现
     *
     * @throws MQClientException
     */
    public static void consumer2() throws MQClientException {
        final MQPullConsumerScheduleService scheduleService = new MQPullConsumerScheduleService("PullConsumerServicel");

        scheduleService.getDefaultMQPullConsumer().setNamesrvAddr("");

        scheduleService.setMessageModel(MessageModel.CLUSTERING);
        scheduleService.registerPullTaskCallback("testPullConsumer", new PullTaskCallback() {
            @Override
            public void doPullTask(MessageQueue mq, PullTaskContext context) {
                MQPullConsumer consumer = context.getPullConsumer();

                try {
                    long offset = consumer.fetchConsumeOffset(mq, false);
                    if (offset < 0) {
                        offset = 0;
                    }

                    PullResult pullResult = consumer.pull(mq, "*", offset, 32);

                    // 根据不同的消息状态做不同的处理
                    switch (pullResult.getPullStatus()) {
                        case FOUND:
                            break;
                        case NO_MATCHED_MSG:
                            break;
                        case NO_NEW_MSG:
                        case OFFSET_ILLEGAL:
                            break;
                        default:
                            break;
                    }


                    consumer.updateConsumeOffset(mq, pullResult.getNextBeginOffset());
                    context.setPullNextDelayTimeMillis(1000);

                } catch (Exception e) {
                    e.printStackTrace();
                }


            }
        });

        scheduleService.start();

    }
}
